import sqlite3
import pandas as pd
from contextlib import contextmanager

class SQLiteMemory:
    def __init__(self, dbpath):  # dbpath：本地数据库文件路径
        self.dbpath = dbpath

    # 连接到本地数据库
    @contextmanager
    def connect(self):
        conn = sqlite3.connect(self.dbpath)
        try:
            # 检查表是否存在，如果不存在则创建
            cursor = conn.cursor()
            cursor.execute("PRAGMA table_info(scan)")
            if not cursor.fetchall():
                cursor.execute("""
                CREATE TABLE scan (
                    item TEXT NOT NULL,
                    sn TEXT NOT NULL,
                    hw TEXT NOT NULL,
                    qty INTEGER NOT NULL,
                    isscan INTEGER DEFAULT 0 CHECK(isscan IN (0, 1))
                )
                """)
                conn.commit()

            yield conn
        finally:
            conn.close()

    # 修改isscan字段
    def update_isscan(self, sql):
        with self.connect() as conn:
            cursor = conn.cursor()
            try:
                cursor.execute(sql)
                conn.commit()
                return True
            except Exception as e:
                conn.rollback()
                return False

    # 初始化本地数据库，把DataFrame插入到本地数据库中
    def initialize_memory_database(self, df, table_name, if_exists='append'):
        with self.connect() as conn:
            try:
                # 将新数据插入到内存数据库的表中
                df.to_sql(table_name, con=conn, if_exists=if_exists, index=False)
                # 提交内存数据库的更改
                conn.commit()
                print("ok")
            except Exception as e:
                print(f"An error occurred: {e}")

    # 使用事务批量删除
    def delete_from_sqlite_with_transaction(self, df, table_name, column):
        """
        使用事务批量删除SQLite数据库中的记录。
        参数：
        df: DataFrame对象，包含要删除的记录的标识。
        table_name: 数据库中的表名。
        column: 要匹配的列名。
        """
        with self.connect() as conn:
            cursor = conn.cursor()
            # 开始一个新的事务
            cursor.execute("BEGIN TRANSACTION;")
            try:
                # 遍历DataFrame对象，删除每一行对应的记录
                for index, row in df.iterrows():
                    # 构造SQL语句
                    sql_delete = f"DELETE FROM {table_name} WHERE {column} = ?"
                    # 执行SQL语句
                    cursor.execute(sql_delete, (row[column],))
                # 提交事务
                conn.commit()
            except Exception as e:
                print(f"An error occurred: {e}")
                # 回滚事务
                conn.rollback()

    # 查询内存数据库中的数据，如果查询成功，返回DataFrame，否则返回None
    def query(self, sql):
        with self.connect() as conn:
            return pd.read_sql_query(sql, conn)

    # 将DataFrame批量导入到数据库中，减少数据库操作次数，可以提高数据导入效率，如果导入成功，返回True，否则返回False
    def insert_dataframe(self, df):
        if df.empty:
            return  # 如果DataFrame为空，则不进行插入操作

        with self.connect() as conn:
            cur = conn.cursor()
            try:
                conn.execute("BEGIN TRANSACTION")  # 开始事务
                sql = 'INSERT INTO scan(item, sn, hw, qty, isscan) VALUES (?, ?, ?, ?, ?)'
                cur.executemany(sql, df.values.tolist())
                conn.commit()

                if cur.rowcount != len(df):  # 检查实际插入的行数是否与DataFrame的行数相等
                    raise RuntimeError("Not all rows were inserted successfully. Rolling back transaction.")
            except Exception as e:
                conn.rollback()  # 回滚事务
                return False

        return True

    # 清空表数据
    def clear_table(self):
        with self.connect() as conn:
            cursor = conn.cursor()
            try:
                cursor.execute("DELETE FROM scan")
                conn.commit()
                return True
            except Exception as e:
                conn.rollback()
                return False

    def update_isscan_batch(self, sql_list):
        """
        批量修改isscan字段。

        参数：
        sql_list: 包含多条修改语句的 SQL 语句列表。
        """
        with self.connect() as conn:
            cursor = conn.cursor()
            cursor.execute("BEGIN TRANSACTION;")
            try:
                for sql in sql_list:
                    cursor.execute(sql)
                    # if cursor.rowcount != 1:  # 检查是否成功修改了一条记录
                    #     raise RuntimeError("Failed to update one or more records. Rolling back transaction.")
                conn.commit()
                return True
            except Exception as e:
                print(f"An error occurred: {e}")
                conn.rollback()
                return False
